package com.leiyk.rateLimiter.aop;

import com.leiyk.rateLimiter.annotation.TokenRateLimiter;
import com.leiyk.rateLimiter.exception.RateLimiterException;
import jakarta.annotation.Resource;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.Signature;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.annotation.Pointcut;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.StringRedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.Arrays;

/**
 * @Description: 令牌桶算法限流注解aop实现
 * @Author: leiyk6
 * @Date: 2024/7/9 14:43
 */
@Component
@Aspect
public class TokenRateLimiterAop {
    public static final String srcPath = "src/main/resources/lua/token_rate.lua";

    @Resource
    StringRedisTemplate stringRedisTemplate;
    @Pointcut("@annotation(com.leiyk.rateLimiter.annotation.TokenRateLimiter)")
    public void pc(){};

    @Around("pc()")
    Object around(ProceedingJoinPoint joinPoint) throws Throwable {
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();
        TokenRateLimiter annotation = method.getAnnotation(TokenRateLimiter.class);
        int rate = annotation.rate();
        int capacity = annotation.capacity();
        int consume = annotation.consume();
        String key = method.toString();
        DefaultRedisScript<Long> script = new DefaultRedisScript<>();
        script.setLocation(new ClassPathResource(srcPath));
        Long res = stringRedisTemplate.execute(script, Arrays.asList(key), rate, capacity, consume);
        if(res != 1L){
            throw RateLimiterException.rejectRequest();
        }
        return joinPoint.proceed();

    }



}
